import wespeaker


class SimModel:
    def __init__(self, gpuid, fs, lan="chinese"):
        if lan == "chinese":
            self.model = wespeaker.load_model("chinese")
        else:
            self.model = wespeaker.load_model("english")
        self.model.set_gpu(gpuid)
        self.fs = fs

    def sim(self, a1, a2):
        """
        :params a1 shape: [batch, time]
        :params a2 shape: [batch, time]
        :return the average similarity score
        """
        sim = 0
        for t1, t2 in zip(a1, a2):
            t1, t2 = t1.unsqueeze(0), t2.unsqueeze(0)
            sim += self.model.compute_similarity_torch(t1, t2, self.fs)
        return sim / len(a1)
